import os
from torchvision import transforms
from PIL import Image


class MyResizeProcessor:
    def __init__(self, TARGET_SIZE=224):
        self.transform = transforms.Compose([
            transforms.Resize((TARGET_SIZE, TARGET_SIZE)),
        ])

    def __call__(self, x):
        return self.transform(x)


def jpg_size_modify(root, path_jpg, name_jpg):
    if not path_jpg.endswith('jpg'):
        print('wrong pic format')
        return None

    img = Image.open(path_jpg).convert('RGB')
    print(type(img), img.size)

    resize_processor = MyResizeProcessor()
    img_resized = resize_processor(img)
    print(type(img_resized), img_resized.size)

    try:
        os.mkdir(os.path.join(root, 'big'))
    except:
        print('result directory already exists')
    finally:
        img_resized.save(os.path.join(root, 'big', '{}_resized.jpg'.format(name_jpg[:-4])))


def find_and_resize_pics(start_root='.'):
    for root, dirs, files in os.walk(start_root):
        # print(root)
        if not root.endswith('big'):
            for file in files:
                if file.endswith('jpg'):
                    path_jpg = os.path.join(root, file)
                    jpg_size_modify(root=root, path_jpg=path_jpg, name_jpg=file)


find_and_resize_pics('./cache/raw/')
print('done')
